package com.ideaaedi.component.compile;

import com.ideaaedi.commonds.clazz.ClassNameGrabber;
import org.apache.commons.io.IOUtils;

import javax.tools.Diagnostic;
import javax.tools.DiagnosticCollector;
import javax.tools.JavaCompiler;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.StandardJavaFileManager;
import javax.tools.ToolProvider;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.UndeclaredThrowableException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;


/**
 * 动态编译器
 *
 * @author <a href="https://github.com/alibaba/arthas">Arthas</a>
 * @author JustryDeng
 * @since 2021/9/20 18:23:21
 */
@SuppressWarnings("all")
public class Compiler {
    private final JavaCompiler javaCompiler = ToolProvider.getSystemJavaCompiler();
    private final StandardJavaFileManager standardFileManager;
    private final List<String> options = new ArrayList<String>();
    private final DynamicClassLoader dynamicClassLoader;

    private final Collection<JavaFileObject> compilationUnits = new ArrayList<JavaFileObject>();
    private final List<Diagnostic<? extends JavaFileObject>> errors = new ArrayList<Diagnostic<? extends JavaFileObject>>();
    private final List<Diagnostic<? extends JavaFileObject>> warnings = new ArrayList<Diagnostic<? extends JavaFileObject>>();

    private Compiler(ClassLoader classLoader) {
        if (javaCompiler == null) {
            throw new IllegalStateException(
                            "Can not load JavaCompiler from javax.tools.ToolProvider#getSystemJavaCompiler(),"
                                            + " please confirm the application running in JDK not JRE.");
        }
        standardFileManager = javaCompiler.getStandardFileManager(null, null, null);

        options.add("-Xlint:unchecked");
        dynamicClassLoader = new DynamicClassLoader(classLoader);
    }

    public static Compiler create() {
        return new Compiler(Compiler.class.getClassLoader());
    }
    
    public static Compiler create(ClassLoader classLoader) {
        return new Compiler(classLoader);
    }

    /**
     * @see Compiler#addSource(String, String)
     */
    public Compiler addSource(String source) {
        addSource(new CharSequenceJavaFileObject(ClassNameGrabber.obtainFromSourceCode(source), source));
        return this;
    }
    

    /**
     * 添加要编译的源码
     *
     * @param classLongName
     *            全类名
     * @param source
     *            java源码
     */
    public Compiler addSource(String classLongName, String source) {
        Objects.requireNonNull(source, "source cannot be null.");
        addSource(new CharSequenceJavaFileObject(classLongName, source));
        return this;
    }
    
    /**
     * @see Compiler#addSource(String, byte[])
     */
    public Compiler addSource(byte[] source) {
        Objects.requireNonNull(source, "source cannot be null.");
        return addSource(ClassNameGrabber.obtainFromSourceCode(new String(source)), source);
    }
    
    /**
     * 添加要编译的源码
     *
     * @param classLongName
     *            全类名
     * @param source
     *            java源码
     */
    public Compiler addSource(String classLongName, byte[] source) {
        Objects.requireNonNull(source, "source cannot be null.");
        addSource(new CharSequenceJavaFileObject(classLongName, new String(source, StandardCharsets.UTF_8)));
        return this;
    }
    
    
    /**
     * @see Compiler#addSource(String, InputStream)
     */
    public Compiler addSource(InputStream source) {
        try {
            Objects.requireNonNull(source, "source cannot be null.");
            String classLongName = ClassNameGrabber.obtainFromSourceCode(IOUtils.toString(source, StandardCharsets.UTF_8));
            return addSource(classLongName, source);
        } catch (IOException e) {
            throw new UndeclaredThrowableException(e);
        }
    }
    
    /**
     * 添加要编译的源码
     *
     * @param classLongName
     *            全类名
     * @param source
     *            java源码
     */
    public Compiler addSource(String classLongName, InputStream source) {
        Objects.requireNonNull(source, "source cannot be null.");
        addSource(new CharSequenceJavaFileObject(classLongName, source));
        return this;
    }

    /**
     * 添加要编译的源码
     *
     * @param javaFileObject
     *            java源码
     */
    private Compiler addSource(JavaFileObject javaFileObject) {
        Objects.requireNonNull(javaFileObject, "javaFileObject cannot be null.");
        compilationUnits.add(javaFileObject);
        return this;
    }
    
    /**
     * 编译并返回 全类名-Class实例 map
     *
     * @return  全类名-Class实例
     */
    public Map<String, Class<?>> build() {

        errors.clear();
        warnings.clear();

        JavaFileManager fileManager = new DynamicJavaFileManager(standardFileManager, dynamicClassLoader);

        DiagnosticCollector<JavaFileObject> collector = new DiagnosticCollector<JavaFileObject>();
        JavaCompiler.CompilationTask task = javaCompiler.getTask(null, fileManager, collector, options, null,
                        compilationUnits);

        try {

            if (!compilationUnits.isEmpty()) {
                boolean result = task.call();

                if (!result || collector.getDiagnostics().size() > 0) {

                    for (Diagnostic<? extends JavaFileObject> diagnostic : collector.getDiagnostics()) {
                        switch (diagnostic.getKind()) {
                        case NOTE:
                        case MANDATORY_WARNING:
                        case WARNING:
                            warnings.add(diagnostic);
                            break;
                        case OTHER:
                        case ERROR:
                        default:
                            errors.add(diagnostic);
                            break;
                        }

                    }

                    if (!errors.isEmpty()) {
                        throw new DynamicCompilerException("Compilation Error", errors);
                    }
                }
            }

            return dynamicClassLoader.getClasses();
        } catch (Throwable e) {
            throw new DynamicCompilerException(e, errors);
        } finally {
            compilationUnits.clear();

        }

    }
    
    /**
     * 编译并返回 全类名-class字节数组 map
     *
     * @return  全类名-class字节数组
     */
    public Map<String, byte[]> buildByteCodes() {

        errors.clear();
        warnings.clear();

        JavaFileManager fileManager = new DynamicJavaFileManager(standardFileManager, dynamicClassLoader);

        DiagnosticCollector<JavaFileObject> collector = new DiagnosticCollector<JavaFileObject>();
        JavaCompiler.CompilationTask task = javaCompiler.getTask(null, fileManager, collector, options, null,
                        compilationUnits);

        try {

            if (!compilationUnits.isEmpty()) {
                boolean result = task.call();

                if (!result || collector.getDiagnostics().size() > 0) {

                    for (Diagnostic<? extends JavaFileObject> diagnostic : collector.getDiagnostics()) {
                        switch (diagnostic.getKind()) {
                        case NOTE:
                        case MANDATORY_WARNING:
                        case WARNING:
                            warnings.add(diagnostic);
                            break;
                        case OTHER:
                        case ERROR:
                        default:
                            errors.add(diagnostic);
                            break;
                        }

                    }

                    if (!errors.isEmpty()) {
                        throw new DynamicCompilerException("Compilation Error", errors);
                    }
                }
            }

            return dynamicClassLoader.getByteCodes();
        } catch (ClassFormatError e) {
            throw new DynamicCompilerException(e, errors);
        } finally {
            compilationUnits.clear();

        }

    }

    private List<String> diagnosticToString(List<Diagnostic<? extends JavaFileObject>> diagnostics) {

        List<String> diagnosticMessages = new ArrayList<String>();

        for (Diagnostic<? extends JavaFileObject> diagnostic : diagnostics) {
            diagnosticMessages.add(
                            "line: " + diagnostic.getLineNumber() + ", message: " + diagnostic.getMessage(Locale.US));
        }

        return diagnosticMessages;

    }

    public List<String> getErrors() {
        return diagnosticToString(errors);
    }

    public List<String> getWarnings() {
        return diagnosticToString(warnings);
    }

    public ClassLoader getClassLoader() {
        return dynamicClassLoader;
    }
}
